import re
import sys
import json
from langgraph.graph import END, StateGraph
from typing import Dict, TypedDict, Literal
from pydantic import BaseModel, Field
from doubao import doubao_chat


# 定义工具参数模型 (自动生成提示词的基础)
class MathOperationParams(BaseModel):
    """数学计算参数模型"""
    a: int = Field(..., title="第一个整数")
    b: int = Field(..., title="第二个整数")
    operation: Literal["add", "subtract"] = Field(..., title="操作类型", description="只能是 add 或 subtract，其中 add 表示加法，subtract 表示减法")


# 定义状态结构
class AgentState(TypedDict):
    user_input: str  # 用户输入的字符串
    params: MathOperationParams  # 解析出的参数
    result: int      # 计算结果


def auto_generate_prompt(schema: type[BaseModel], user_input: str) -> str:
    """
    根据工具参数schema自动生成提示词 (兼容Python 3.12)
    """
    # 自动生成提示词
    prompt = f"""
    请从以下用户查询中提取结构化信息：
    "{user_input}"

    你需要提取以下字段，用json schema描述如下：
    {schema.model_json_schema()}

    输出要求：
    1. 必须是纯JSON格式，仅包含上述字段
    2. 不要包含任何额外说明或文本
    3. 如果无法确定某个字段的值，请使用null

    示例输出：
    {json.dumps({"a": 5, "b": 3, "operation": "add"}, ensure_ascii=False)}
    """
    return prompt.strip()


def llm_chat(state: AgentState) -> Dict:
    """使用豆包大模型解析用户输入，提取数字和操作类型"""
    print(f"\n🧠 用户输入: '{state['user_input']}'")

    # 构建给大模型的提示
    prompt = auto_generate_prompt(MathOperationParams, state["user_input"])
    # print(prompt)

    # 调用豆包大模型API
    llm_output = doubao_chat(api_key=sys.argv[1], text=prompt)

    # 提取模型返回的内容
    content = llm_output.message.content
    print(f"🤖 豆包大模型原始输出: {content}")
    if not content:
        raise ValueError("模型返回内容为空")

    # 从输出中提取JSON部分（模型可能会添加额外说明）
    json_match = re.search(r'\{.*\}', content, re.DOTALL)
    if not json_match:
        raise ValueError("未找到有效的JSON输出")

    # 将JSON字符串转换为字典
    result_dict = json.loads(json_match.group(0))

    # 验证并转换参数
    validated_params = MathOperationParams(**result_dict)
    print(f"✅ 验证通过: {validated_params}")

    # 验证结果
    return {"params": validated_params}


# 计算节点函数
def add_numbers(state: AgentState) -> Dict[str, int]:
    params = state["params"]
    result = params.a + params.b
    print(f"➕ 计算结果: {params.a} + {params.b} = {result}")
    return {"result": result}


def subtract_numbers(state: AgentState) -> Dict[str, int]:
    params = state["params"]
    result = params.a - params.b
    print(f"➖ 计算结果: {params.a} - {params.b} = {result}")
    return {"result": result}


# 结果展示节点
def display_result(state: AgentState) -> Dict:
    print(f"✅ 最终结果: {state['result']}")
    return {}


# 构建图工作流
def create_workflow():
    # 初始化状态图
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("chat", llm_chat)
    workflow.add_node("add", add_numbers)
    workflow.add_node("subtract", subtract_numbers)
    workflow.add_node("display", display_result)

    # 设置入口点
    workflow.set_entry_point("chat")

    # 根据解析结果路由
    workflow.add_conditional_edges(
        "chat",
        lambda state: state["params"].operation,
        {
            "add": "add",
            "subtract": "subtract",
        }
    )

    # 连接计算节点到显示节点
    workflow.add_edge("add", "display")
    workflow.add_edge("subtract", "display")

    # 连接到结束
    workflow.add_edge("display", END)

    return workflow.compile()


# 执行示例
if __name__ == "__main__":
    # 初始化工作流
    app = create_workflow()

    # 测试案例
    test_cases = [
        "10减去7",
        "28加90等于多少？",
        "计算20和15的和",
        "100 - 50",
        "8和4的差是多少？",
        "简单的计算：3和5",  # 测试默认操作
        "只有数字：42 100",  # 测试无操作符
    ]

    for input_str in test_cases:
        print("\n" + "="*50)
        print(f"测试输入: '{input_str}'")
        result = app.invoke({"user_input": input_str})
        print(f"完整输出: {result}")
        break
